#!/usr/bin/env python
# -*- encoding: utf-8 -*-


import time
import json

from elasticsearch import Elasticsearch
from pyspider.database.base.errordb import ErrorDB as BaseErrorDB


class ErrorDB(BaseErrorDB):
    __type__ = 'error'

    def __init__(self, hosts, index='pyspider'):
        self.index = index
        self._changed = False
        self.es = Elasticsearch(hosts=hosts)

        self.es.indices.create(index=self.index, ignore=400)
        if not self.es.indices.get_mapping(index=self.index, doc_type=self.__type__):
            self.es.indices.put_mapping(index=self.index, doc_type=self.__type__, body={
                '_all': {'enabled': False},
                'properties': {
                    'mark': {'type': 'string'},
                    'errortype': {'type': 'integer'}
                }
            })

    @staticmethod
    def _parse(data):
        if not data:
            return data
        for each in ('schedule', 'fetch', 'process', 'track', 'mark'):
            data[each] = json.loads(data[each]) if data.get(each, None) else {}
        return data

    @staticmethod
    def _stringify(data):
        for each in ('schedule', 'fetch', 'process', 'track', 'mark'):
            if each in data:
                data[each] = json.dumps(data[each])
        return data

    def get_tasks(self, error_type, mark=False, limit=1000):
        self.refresh()
        return [self._parse(record['_source']) for record in self.es.search(
            index=self.index, doc_type=self.__type__,
            body={'query': {'bool': {
             'must': [{'term': {'mark': mark}}, {'term': {'errortype': error_type}}],
             'minimum_should_match': 1
            }}}, size=limit).get('hits', {}).get('hits', [])]

    def get_task(self, project, taskid):
        if self._changed:
            self.refresh()
        ret = self.es.get(index=self.index, doc_type=self.__type__, id="%s:%s" % (project, taskid), ignore=404)
        return self._parse(ret.get('_source', None))

    def count(self, error_type, hour=None):
        self.refresh()
        if hour:
            ret = self.es.search(
                index=self.index, doc_type=self.__type__,
                body={'query': {'bool': {'must': [
                    {'term': {'mark': False}}, {'term': {'errortype': error_type}},
                    {'range': {'updatetime': {'gt': time.time() - hour * 3600}}}]}},
                    'aggs': {'status': {'terms': {'field': 'mark'}}}
                }, _source=False)
        else:
            ret = self.es.search(index=self.index, doc_type=self.__type__,
                                 body={'query': {'bool': {'must': [
                                     {'term': {'mark': False}}, {'term': {'errortype': error_type}}]}},
                                     'aggs': {'status': {'terms': {'field': 'mark'}}}}, _source=False)
        for each in ret['aggregations']['status']['buckets']:
            return each['doc_count']
        return 0

    def count_by_date(self, error_type, date):
        return self.es.count(index=self.index, doc_type=self.__type__,
                             body={'query': {'bool': {'must': [{'term': {'errortype': error_type}},
                                                               {'match_phrase': {'date': date}}]}}}
                             ).get('count', 0)

    def insert(self, project, taskid, obj):
        self._changed = True
        obj = dict(obj)
        obj['taskid'] = taskid
        obj['project'] = project
        obj['updatetime'] = time.time()
        obj['mark'] = False
        return self.es.index(index=self.index, doc_type=self.__type__,
                             body=self._stringify(obj), id='%s:%s' % (project, taskid))

    def insert_vacancy(self, error_type, project, date, group):
        obj = {
            'errortype': error_type,
            'project': project,
            'date': date,
            'file': file,
            'group': group,
        }
        return self.es.index(index=self.index, doc_type=self.__type__,
                             body=obj, id='%s:%s' % (project, date))

    def select(self, error_type, hour=None, fields=None, offset=0, limit=1000):
        if not hour:
            for record in self.es.search(index=self.index, doc_type=self.__type__,
                                         body={'query': {'term': {'errortype': error_type}}},
                                         _source_include=fields or [], from_=offset, size=limit,
                                         sort="updatetime:desc"
                                         ).get('hits', {}).get('hits', []):
                yield record['_source']
        else:
            for record in self.es.search(index=self.index, doc_type=self.__type__,
                                         body={'query': {'bool': {'must': [
                                             {'term': {'errortype': error_type}},
                                             {'range': {'updatetime': {'gt': time.time() - hour * 3600}}}]}}},
                                         _source_include=fields or [], from_=offset, size=limit,
                                         sort="updatetime:desc"
                                         ).get('hits', {}).get('hits', []):
                yield record['_source']

    def update(self, project, taskid, obj):
        self._changed = True
        obj = dict(obj)
        return self.es.update(index=self.index, doc_type=self.__type__, id='%s:%s' % (project, taskid),
                              body={'doc': self._stringify(obj)}, ignore=404)

    def refresh(self):
        if self._changed:
            self._changed = False
            self.es.indices.refresh(index=self.index)

    def delete(self, project, date, refresh=True):
        self.refresh()
        self.es.delete(index=self.index, doc_type=self.__type__, id="%s:%s" % (project, date))
        self.refresh()
